package maps;

import java.util.Comparator;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;

public class MyTreeMap<K, V> implements MyMap<K, V> {

	private final Comparator<K> comparator;
	private int size;
	private SimpleEntry<K, V> root;

	public MyTreeMap() {
		this.comparator = null;

	}

	public MyTreeMap(Comparator<K> comparator) {
		this.comparator = comparator;

	}

	@Override
	public void clear() {
		root = null;
		size = 0;

	}

	@Override
	public boolean containsKey(K key) {

		return (get(key) != null);
	}

	@Override
	public boolean containsValue(V value) {
		if (isEmpty())
			throw new NullPointerException();
		List<K> keys = getKeysList();
		for (int i = 0; i < keys.size(); i++) {
			V val = get(keys.get(i));
			if (val.equals(value))
				return true;
		}
		return false;
	}

	@Override
	public V get(K key) {
		if (isEmpty() || key == null) {
			throw new NullPointerException();
		}
		SimpleEntry<K, V> x = root;

		if (comparator == null) {
			while (x != null) {
				@SuppressWarnings("unchecked")
				int cmp = ((Comparable<K>) key).compareTo(x.key);
				if (cmp < 0) {
					x = x.left;
				} else {
					if (cmp > 0) {
						x = x.right;
					} else {
						return x.value;
					}
				}
			}
		} else {
			while (x != null) {

				int cmp = comparator.compare(key, x.key);
				if (cmp < 0) {
					x = x.left;
				} else {
					if (cmp > 0) {
						x = x.right;
					} else {
						return x.value;
					}
				}
			}
		}
		return null;
	}

	@Override
	public boolean isEmpty() {

		return root == null;
	}

	@Override
	public V put(K key, V value) {
		if (key == null)
			throw new NullPointerException();
		root = put(root, key, value);
		size++;
		return null;
	}

	@Override
	public V remove(K key) {
		delete(key);
		size--;
		return null;
	}

	private SimpleEntry<K, V> deleteMin(SimpleEntry<K, V> h) {
		if (h.left == null)
			return null;

		if (!isRed(h.left) && !isRed(h.left.left))
			h = moveRedLeft(h);

		h.left = deleteMin(h.left);
		return balance(h);
	}

	private SimpleEntry<K, V> delete(SimpleEntry<K, V> x, K key) {
		if (x == null)
			return null;
		if (comparator == null) {
			@SuppressWarnings("unchecked")
			int cmp = ((Comparable<K>) key).compareTo(x.key);
			if (cmp < 0)
				x.left = delete(x.left, key);
			else if (cmp > 0)
				x.right = delete(x.right, key);
			else {
				if (x.right == null)
					return x.left;
				if (x.left == null)
					return x.right;
				SimpleEntry<K, V> t = x;
				x = min(t.right);
				x.right = deleteMin(t.right);
				x.left = t.left;
			}
		} else {
			int cmp = comparator.compare(key, x.key);
			if (cmp < 0)
				x.left = delete(x.left, key);
			else if (cmp > 0)
				x.right = delete(x.right, key);
			else {
				if (x.right == null)
					return x.left;
				if (x.left == null)
					return x.right;
				SimpleEntry<K, V> t = x;
				x = min(t.right);
				x.right = deleteMin(t.right);
				x.left = t.left;
			}
		}
		return balance(x);
	}

	private void delete(K key) {
		if (!containsKey(key)) {
			System.err.println("symbol table does not contain " + key);
			throw new NullPointerException();

		}

		if (!isRed(root.left) && !isRed(root.right))
			root.color = RED;

		root = delete(root, key);
		if (!isEmpty())
			root.color = BLACK;

	}

	private SimpleEntry<K, V> min(SimpleEntry<K, V> x) {
		assert x != null;
		if (x.left == null)
			return x;
		else
			return min(x.left);
	}

	private SimpleEntry<K, V> moveRedLeft(SimpleEntry<K, V> h) {
		assert (h != null);
		assert isRed(h) && !isRed(h.left) && !isRed(h.left.left);

		flipColors(h);
		if (isRed(h.right.left)) {
			h.right = rotateRight(h.right);
			h = rotateLeft(h);
		}
		return h;
	}

	private SimpleEntry<K, V> balance(SimpleEntry<K, V> h) {
		assert (h != null);

		if (isRed(h.right))
			h = rotateLeft(h);
		if (isRed(h.left) && isRed(h.left.left))
			h = rotateRight(h);
		if (isRed(h.left) && isRed(h.right))
			flipColors(h);

		return h;
	}

	@Override
	public int size() {

		return size;
	}

	@Override
	public Iterator<K> entryIterator() {

		return new Iterator<K>() {
			List<K> sortedList = getKeysList();
			Iterator<K> iter = sortedList.iterator();

			@Override
			public boolean hasNext() {

				return iter.hasNext();
			}

			@Override
			public K next() {

				return iter.next();
			}

			@Override
			public void remove() {
				iter.remove();

			}
		};
	}

	public List<K> getKeysList() {
		List<K> q = new LinkedList<K>();
		inorder(root, q);
		return q;

	}

	private void inorder(SimpleEntry<K, V> x, List<K> q) {
		if (x == null)
			return;
		inorder(x.left, q);
		q.add(x.key);
		inorder(x.right, q);
	}

	private boolean isRed(SimpleEntry<K, V> x) {
		if (x == null)
			return false;
		return x.color == RED;
	}

	private SimpleEntry<K, V> rotateLeft(SimpleEntry<K, V> h) {

		assert isRed(h.right);
		SimpleEntry<K, V> x = h.right;
		h.right = x.left;
		x.left = h;
		x.color = h.color;
		h.color = RED;
		return x;
	}

	private SimpleEntry<K, V> rotateRight(SimpleEntry<K, V> h) {

		assert isRed(h.left);
		SimpleEntry<K, V> x = h.left;
		h.left = x.right;
		x.right = h;
		x.color = h.color;
		h.color = RED;
		return x;
	}

	private void flipColors(SimpleEntry<K, V> h) {
		assert !isRed(h);
		assert isRed(h.left);
		assert isRed(h.right);

		h.color = RED;
		h.left.color = BLACK;
		h.right.color = BLACK;
	}

	private SimpleEntry<K, V> put(SimpleEntry<K, V> h, K key, V value) {
		if (h == null)
			return new SimpleEntry<K, V>(key, value, RED);
		if (comparator == null) {
			@SuppressWarnings("unchecked")
			int cmp = ((Comparable<K>) key).compareTo(h.key);
			if (cmp < 0) {
				h.left = put(h.left, key, value);
			} else {
				if (cmp > 0) {
					h.right = put(h.right, key, value);
				} else {

					h.value = value;
				}
			}
		} else {
			int cmp = comparator.compare(key, h.key);
			if (cmp < 0) {
				h.left = put(h.left, key, value);
			} else {
				if (cmp > 0) {
					h.right = put(h.right, key, value);
				} else {
					h.value = value;
				}
			}
		}

		if (isRed(h.right) && !isRed(h.left))
			h = rotateLeft(h);
		if (isRed(h.left) && isRed(h.left.left))
			h = rotateRight(h);
		if (isRed(h.left) && isRed(h.right))
			flipColors(h);

		return h;

	}

	private static final boolean RED = true;
	private static final boolean BLACK = false;

	static class SimpleEntry<K, V> implements MyMap.Entry<K, V> {

		K key;
		V value;
		SimpleEntry<K, V> left;
		SimpleEntry<K, V> right;
		boolean color;

		public SimpleEntry(K key, V value, boolean color) {
			this.key = key;
			this.value = value;
			this.color = color;

		}

		@Override
		public K getKey() {

			return key;
		}

		@Override
		public V getValue() {

			return value;
		}

		@Override
		public Object setValue(V value) {
			V oldValue = this.value;
			this.value = value;
			return oldValue;
		}

	}

}
